#include "structMeshDistribution.hpp"
#include "structMeshPartitioning.hpp"
#include "structMeshOverlap.hpp"
#include "util.hpp"

namespace hms {


/** @todo put this into structMeshFunctions.hpp and see where it can be used */
Index firstEdge( BoundaryLocation loc, const Ref<const Array2i>& meshSize ){
	switch(loc){
		case BoundaryLocation::bottom:
			return 0;
		case BoundaryLocation::left:
			/** @todo
			 * unify this number as function call nHorizontalEdges(meshSize) */
			return nHorizontalEdges(meshSize);
		case BoundaryLocation::right:
			return meshSize.x() * (meshSize.y() + 2);
		case BoundaryLocation::top:
			return meshSize.prod();
		default:
			UNREACHABLE;
	}
}

Index lastEdge( BoundaryLocation loc, const Ref<const Array2i>& meshSize ){
	switch(loc){
		case BoundaryLocation::bottom:
			return meshSize.x() - 1;
		case BoundaryLocation::left:
			return 2 * meshSize.prod() + meshSize.y() - 1;
		case BoundaryLocation::right:
			return 2 * meshSize.prod() + meshSize.sum() - 1;
		case BoundaryLocation::top:
			return meshSize.prod() + meshSize.x() - 1;
		default:
			UNREACHABLE;
	}
}



UniMesh distributeMesh( 
	const UniMesh& mesh,
	const MPI_Comm_H& comm,
	const Ref<const Array2i>&,
	const Ref<const Array2i>& meshSizeLocal,
	const Ref<const Array2i>&,
	Index ng
){
	Array2s diff;
	Array2s origin;
	if ( comm.isMaster() ){
		diff = mesh.diff();
		origin = mesh.origin();
	}
	MPI_Bcast_H(diff, comm, comm.master(), false);
	MPI_Bcast_H(origin, comm, comm.master(), false);

	return { meshSizeLocal, std::move(diff), std::move(origin), ng };
}


RectMesh distributeMesh( 
	const RectMesh& mesh,
	const MPI_Comm_H& comm,
	const Ref<const Array2i>&,
	const Ref<const Array2i>& meshSizeLocal,
	const Ref<const Array2i>& offset,
	Index ng
){
	auto getCoords = [&]( const ArrayXs& globalCoords, int dim ){

		MPI_Win_H win { const_cast<ArrayXs&>( globalCoords ), comm };

		ArrayXs coords ( 2*ng + meshSizeLocal[dim] + 1 );

		win.fence();
		if ( !comm.isMaster() )
			win.get( coords, 0, coords.size(), comm.master(), offset[dim] );
		else
			coords = globalCoords.head( coords.size() );
		win.fence();

		return coords;
	};

	ArrayXs xCoords { getCoords(mesh.xCoords(), 0) };
	ArrayXs yCoords { getCoords(mesh.yCoords(), 1) };

	return { std::move(xCoords), std::move(yCoords), ng, true };
}


StructMesh distributeMesh( 
	const StructMesh& mesh,
	const MPI_Comm_H& comm,
	const Ref<const Array2i>& meshSizeGlobal,
	const Ref<const Array2i>& meshSize,
	const Ref<const Array2i>& offset,
	Index ng
){
	Matrix2Xs vertices = decomposeArray(
		mesh.vertices(), meshSizeGlobal + 1, meshSize + 1, offset,
		comm, comm.master()
	);

	constexpr bool recalcValues { true };

	if constexpr ( recalcValues ){

	VectorXs edgeLengths ( 2 * meshSize.prod() + meshSize.sum() );
	Matrix2Xs edgeNormals ( 2, edgeLengths.size() );

	Array2i hEdgesSizeLocal { meshSize.x(), meshSize.y() + 1 };
	Array2i hEdgesSizeGlobal { meshSizeGlobal.x(), meshSizeGlobal.y() + 1 };
	Array2i vEdgesSizeLocal { meshSize.x() + 1, meshSize.y() };
	Array2i vEdgesSizeGlobal { meshSizeGlobal.x() + 1, meshSizeGlobal.y() };

	auto hEdgeLengths = edgeLengths.head( hEdgesSizeLocal.prod() );
	auto vEdgeLengths = edgeLengths.tail( vEdgesSizeLocal.prod() );

	auto hEdgeNormals = edgeNormals.leftCols( hEdgesSizeLocal.prod() );
	auto vEdgeNormals = edgeNormals.rightCols( vEdgesSizeLocal.prod() );

	decomposeArray(
		mesh.edges().lengths().head(
			comm.isMaster() ? hEdgesSizeGlobal.prod() : 0 ),
		hEdgeLengths, 1,
		hEdgesSizeGlobal, hEdgesSizeLocal, offset, comm, comm.master()
	);
	decomposeArray(
		mesh.edges().lengths().tail(
			comm.isMaster() ? vEdgesSizeGlobal.prod() : 0 ),
		vEdgeLengths, 1,
		vEdgesSizeGlobal, vEdgesSizeLocal, offset, comm, comm.master()
	);

	decomposeArray(
		mesh.edges().normals().leftCols(
			comm.isMaster() ? hEdgesSizeGlobal.prod() : 0 ),
		hEdgeNormals, 2,
		hEdgesSizeGlobal, hEdgesSizeLocal, offset, comm, comm.master()
	);
	decomposeArray(
		mesh.edges().normals().rightCols(
			comm.isMaster() ? vEdgesSizeGlobal.prod() : 0 ),
		vEdgeNormals, 2,
		vEdgesSizeGlobal, vEdgesSizeLocal, offset, comm, comm.master()
	);

	VectorXs cellAreas = decomposeArray(
		mesh.cells().areas(), meshSizeGlobal + 2*ng, meshSize + 2*ng, offset,
		comm, comm.master()
	);

	Matrix2Xs cellCentroids = decomposeArray(
		mesh.cells().centroids(), meshSizeGlobal + 2*ng, meshSize+2*ng, offset,
		comm, comm.master()
	);

	return {
		std::move(vertices), meshSize, ng,
		std::move(edgeLengths), std::move(edgeNormals),
		std::move(cellAreas), std::move(cellCentroids)
	};

	} else {
		return {
			std::move(vertices), meshSize, ng
		};
	}
}



void redistributeMesh( 
	UniMesh& mesh,
	const MPI_Comm_H&,
	const Ref<const Array2i>&,
	const Ref<const Array2i>&,
	const Ref<const Array2i>& meshSizeLocal,
	const Ref<const Array2i>&,
	const Ref<const Array2Xi>&,
	const Ref<const Array2Xi>&
){
	mesh = { meshSizeLocal, mesh.diff(), mesh.origin(), mesh.nGhostLayers() };
}


void redistributeMesh(
	RectMesh& mesh,
	const MPI_Comm_H& comm,
	const Ref<const Array2i>& part,
	const Ref<const Array2i>& nParts,
	const Ref<const Array2i>& meshSize,
	const Ref<const Array2i>& offset,
	const Ref<const Array2Xi>& allSizes,
	const Ref<const Array2Xi>& allOffsets
){
	static constexpr Index ng {1};

	/* easier format for offsets in both dimensions */
	std::array<ArrayXi,2> offsets {
		allOffsets.row(0).reshaped( nParts.x(), nParts.y() )
			.col( part.y() ),
		allOffsets.row(1).reshaped( nParts.x(), nParts.y() )
			.row(0)
	};
	std::array<ArrayXi,2> sizes {
		allSizes.row(0).reshaped( nParts.x(), nParts.y() )
			.col( part.y() ),
		allSizes.row(1).reshaped( nParts.x(), nParts.y() )
			.row(0)
	};

	std::array oldCoords { mesh.xCoords(), mesh.yCoords() };

	auto getCoords = [&]( int dim ) -> ArrayXs {

		std::vector<OverlapSegment> segments;

		Index cur {0};
		int firstPart { findPart(offset[dim], offsets[dim]) };
		for ( int i{firstPart}; i<nParts[dim]; ++i ){
			/* Coordinates do not become invalid after one time step, as ghost
			 * values do. Therefore, the overlap can include ghost cells.
			 * Additionally, we add 1 to the size, because there is always one
			 * more coordinate per dimension than cells.
			 * We add the running counter 'cur' to the lower bound to exclude
			 * those coordinates from the overlap that have already been 
			 * covered.
			 * */
			if ( auto limits = overlap(
				offset(dim) + cur, offset(dim) + meshSize(dim) + 2*ng + 1,
				offsets[dim](i), offsets[dim](i) + sizes[dim](i) + 2*ng + 1
			)){
				int rank { static_cast<int>( dim == 0 ?
					part.y() * nParts.x() + i :
					part.x() + nParts.x() * i )
				};
				segments.push_back( {
					rank,
					limits.value()[1] - limits.value()[0],
					cur + offset(dim) - allOffsets(dim, rank),
					cur
				} );
				cur += segments.back().size;
			}
			/* we've started at the first part with an overlap, so no overlap
			 * means we're beyond the executing rank's coordinates */
			else
				break;
		}

		ArrayXs coords ( meshSize(dim) + 1 + 2*ng );

		MPI_Win_H win { const_cast<ArrayXs&>( oldCoords[dim] ), comm };

		win.fence();
		for ( const OverlapSegment& s : segments ){
			if ( s.rank != comm.rank() )
				win.get( coords, s.ownBeg, s.size, s.rank, s.beg );
			else
				coords.segment( s.ownBeg, s.size ) =
					oldCoords[dim].segment( s.beg, s.size );
		}
		win.fence();

		return coords;
	};
	
	ArrayXs xCoords { getCoords(0) };
	ArrayXs yCoords { getCoords(1) };

	mesh = { std::move(xCoords), std::move(yCoords), ng, true };
}


void redistributeMesh(
	StructMesh& mesh,
	const MPI_Comm_H& comm,
	const Ref<const Array2i>& part,
	const Ref<const Array2i>& nParts,
	const Ref<const Array2i>& meshSize,
	const Ref<const Array2i>& offset,
	const Ref<const Array2Xi>& allSizes,
	const Ref<const Array2Xi>& allOffsets
){
	assert( part.y()*nParts.x() + part.x() == comm.rank() &&
		"Rank and position in global mesh do not match!" );
	assert( (allOffsets == meshPartSizesToOffsets(nParts, allSizes)).all() &&
		"Offsets do not match sizes!" );

	static constexpr Index ng {1};

	std::vector<OverlapBlock> blocks { overlappingBlocks(
		part, nParts, offset, meshSize + 1, allOffsets, allSizes + 1, 0
	) };

	Matrix2Xs vertices = redistributeValues( mesh.vertices(),
		meshSize + 1, offset, 0, allSizes + 1, allOffsets, blocks, comm
	);

	VectorXs edgeLengths ( 2 * meshSize.prod() + meshSize.sum() );
	Matrix2Xs edgeNormals ( 2, edgeLengths.size() );

	Array2i pseudoSize;
	Array2Xi pseudoSizes { allSizes };

	auto getEdgeData = [&]( bool horizontal ){
		
		blocks = overlappingBlocks(
			part, nParts, offset, pseudoSize, allOffsets, pseudoSizes, 0
		);

		Index
			nOwn { pseudoSize.prod() },
			nTarget { pseudoSizes.col( comm.rank() ).prod() };

		redistributeValues(
			(horizontal ?
				edgeLengths.head(nOwn) :
				edgeLengths.tail(nOwn) ).transpose(),
			(horizontal ?
				mesh.edges().lengths().head(nTarget) :
				mesh.edges().lengths().tail(nTarget) ).transpose(),
			pseudoSize, offset, 0, pseudoSizes, allOffsets, blocks, comm
		);

		redistributeValues(
			( horizontal ?
				edgeNormals.leftCols(nOwn) : edgeNormals.rightCols(nOwn) ),
			( horizontal ?
				mesh.edges().normals().leftCols(nTarget) :
				mesh.edges().normals().rightCols(nTarget) ),
			pseudoSize, offset, 0, pseudoSizes, allOffsets, blocks, comm
		);
	};

	/* sizes for horizontal edges */
	pseudoSize = { meshSize.x(), meshSize.y() + 1 };
	pseudoSizes.row(1) += 1;

	getEdgeData(true);

	/* number of vertical edges is analogue to number of horizontal edges */
	pseudoSize = { meshSize.x() + 1, meshSize.y() };
	pseudoSizes.colwise() += Array2i{ 1, -1 };

	getEdgeData(false);

	blocks = overlappingBlocks(
		part, nParts, offset, meshSize, allOffsets, allSizes, ng
	);

	VectorXs cellAreas ( (meshSize + 2*ng).prod() );

	redistributeValues(
		cellAreas.transpose(),
		mesh.cells().areas().transpose(),
		meshSize, offset, ng, allSizes, allOffsets, blocks, comm
	);

	Matrix2Xs cellCentroids = redistributeValues(
		mesh.cells().centroids(),
		meshSize, offset, ng, allSizes, allOffsets, blocks, comm
	);

	mesh = {
		std::move(vertices), std::move(meshSize), ng,
		std::move(edgeLengths), std::move(edgeNormals),
		std::move(cellAreas), std::move(cellCentroids)
	};
}



}